import argparse, os, yaml
import matplotlib.pyplot as plt
import numpy as np

parser = argparse.ArgumentParser(description='Arguments for the program')
parser.add_argument('--category', default="mug", type=str)
parser.add_argument('--grasp_id', default="3", type=str)
args = parser.parse_args()

with open('cfg/class_cfg.yml','r') as stream: 
        class_cfg = yaml.safe_load(stream)
inst_num = float(class_cfg[args.category]['eval_inst_num'])

LEGEND_DICT = {"allfc":'L_transfer + L_collision + L_contact + L_selfcolli + L_fc',
               "all":'L_transfer + L_collision + L_contact + L_selfcolli',
               "awsc":'L_transfer + L_collision + L_contact', 
               "tf":'L_transfer + L_collision', 
               "ct":'L_transfer + L_contact',
               "tt":'L_contact + L_collision',
               "fc": 'L_fc', 
               "trans":'L_transfer', 
               "touch":'L_contact',
               "colli":'L_collision',
               "dm":'no refine',
               }

legend_list = []
x = range(90,-1,-1)
print(x[0], x[5], x[10], x[15], x[-1])
plt.figure()
for exp in ["allfc", "allself", "all","awsc", "tf", "ct", "tt", "fc", "trans", "touch", "colli", "dm"]:
    file_path = 'results/{}/{}/{}.npz'.format(args.category, exp, args.grasp_id)
    if not os.path.exists(file_path):
        continue
    legend_list.append(LEGEND_DICT[exp])
    nums = np.load(file_path)['success_num']
    print("{:7}: {}".format(exp, list([nums[0], nums[5], nums[10], nums[15], nums[-1]])))
    nums = nums / inst_num
    plt.plot(x, nums)

plt.xlabel('rotation threshold in degree')
plt.ylabel('grasp success rate')
plt.legend(legend_list, loc='lower right')
plt.savefig('results/{}/grasp_id_{}.png'.format(args.category, args.grasp_id))
# plt.show()
